import numpy as np
from math import sqrt
import math

# load eval_model_n.json and convert to a float list
def load_performance(pth):
    evaluation = None
    with open(pth, "r") as f:
        for line in f:
            evaluation = line
    evaluation = eval(evaluation)
    evaluation = np.array(evaluation)
    return evaluation

evals = []

for g in range(0,41):
    eval_ = load_performance("./eval_model_"+str(g)+".json")
    evals.append(eval_)

evals_ = []

for e in evals:
    e_ = []
    for ee in e:
        e_.append(ee[1][0])
    evals_.append(e_)

for e in evals_:
    print(e)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

X = np.arange(0, 41, 1)
Y = np.arange(0, 41, 1)
X, Y = np.meshgrid(X, Y)
Z = np.array(evals_)
fig = plt.figure()
ax = Axes3D(fig)
# set Z axis range to 0-1
ax.set_zlim(0, 1)

# label x axis as "generation of environments"
ax.set_xlabel("generation of environments")
# label y axis as "generation of models"
ax.set_ylabel("generation of models")
# label z axis as "success rate"
ax.set_zlabel("success rate")

ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap='spring')
plt.show()
